"""
突破策略

基于价格突破一定周期的最高价或最低价进行交易。

作者: AI Assistant
版本: 1.0.0
日期: 2025-01-06
"""

import pandas as pd
from backtest.strategy import StrategyBase


class BreakoutStrategy(StrategyBase):
    """
    突破策略
    
    当价格突破N日最高价时买入，跌破N日最低价时卖出。
    
    参数:
        period: 突破周期，默认20天
        
    信号规则:
        - 收盘价突破N日最高价：买入
        - 收盘价跌破N日最低价：卖出
        - 其他情况：持有
        
    示例:
        >>> strategy = BreakoutStrategy(period=20)
        >>> engine = BacktestEngine(strategy=strategy, data=data)
        >>> result = engine.run()
    """
    
    def __init__(self, period: int = 20):
        """
        初始化突破策略
        
        参数:
            period: 突破周期（天数）
        """
        if period < 2:
            raise ValueError("突破周期至少为2天")
        
        super().__init__(
            name=f"突破策略({period}日)",
            min_period=period + 1,  # 需要period+1天的数据
            period=period
        )
        
        self.period = period
        self.last_signal = None
    
    def generate_signal(self, data: pd.DataFrame, current_bar: pd.Series) -> str:
        """
        生成交易信号
        
        参数:
            data: 当前及之前的所有数据
            current_bar: 当前K线数据
            
        返回:
            str: 'BUY', 'SELL', 或 'HOLD'
        """
        # 确保有足够的数据
        if len(data) < self.period + 1:
            return 'HOLD'
        
        # 获取前N天的数据（不包括今天）
        historical_data = data.iloc[-(self.period + 1):-1]
        
        # 计算N日最高价和最低价
        highest_high = historical_data['最高'].max()
        lowest_low = historical_data['最低'].min()
        
        # 当前收盘价
        current_close = current_bar['收盘']
        
        # 向上突破信号
        if current_close > highest_high:
            if self.last_signal != 'BUY':
                self.last_signal = 'BUY'
                return 'BUY'
        
        # 向下突破信号
        elif current_close < lowest_low:
            if self.last_signal != 'SELL':
                self.last_signal = 'SELL'
                return 'SELL'
        
        return 'HOLD'
    
    def get_indicator_values(self, data: pd.DataFrame) -> pd.DataFrame:
        """
        获取策略指标值（用于绘图）
        
        参数:
            data: 历史数据
            
        返回:
            pd.DataFrame: 包含最高价和最低价通道的数据框
        """
        result = data.copy()
        result[f'{self.period}日最高'] = data['最高'].rolling(
            window=self.period,
            min_periods=self.period
        ).max()
        result[f'{self.period}日最低'] = data['最低'].rolling(
            window=self.period,
            min_periods=self.period
        ).min()
        return result

